import openai
import os
import json
from math import *


openai.api_key = os.environ["OPENAI_API_KEY_ZHIHU"]
openai.api_base = os.environ["OPENAI_API_BASE_ZHIHU"]
model = "gpt-3.5-turbo"




def print_json(data):
    """
    打印参数。如果参数是有结构的（如字典或列表），则以格式化的 JSON 形式打印；
    否则，直接打印该值。
    """
    if hasattr(data, 'model_dump_json'):
        data = json.loads(data.model_dump_json())

    if (isinstance(data, (list, dict))):
        print(json.dumps(
            data,
            indent=4,
            ensure_ascii=False
        ))
    else:
        print(data)

def get_completion(messages, model="gpt-3.5-turbo"):
    response = openai.ChatCompletion.create(
        model=model,
        messages=messages,
        temperature=0.7,  # 模型输出的随机性，0 表示随机性最小
        tools=[{  # 用 JSON 描述函数。可以定义多个。由大模型决定调用谁。也可能都不调用
            "type": "function",
            "function": {
                "name": "sum",
                "description": "加法器，计算一组数的和",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "numbers": {
                            "type": "array",
                            "items": {
                                "type": "number"
                            }
                        }
                    }
                }
            }
        }],
        stream=True
    )
    return response

def sum123(numbers):
    return sum(numbers)

def getStreamResponse(response):
    function_name, tool_call_id, args, text ,return_delta = "", "", "", "",None

    print("====Streaming====")
    # 需要把 stream 里的 token 拼起来，才能得到完整的 call
    for msg in response:
        delta = msg.choices[0].delta
        if return_delta == None:
            return_delta = delta
        if "tool_calls" in delta and delta.tool_calls:
            if not function_name:
                function_name = delta.tool_calls[0].function.name
                tool_call_id = delta.tool_calls[0].id
                print(function_name)
            args_delta = delta.tool_calls[0].function.arguments
            print(args_delta)  # 打印每次得到的数据
            args = args + args_delta
        elif "content" in delta and delta.content:
            text_delta = delta.content
            print(text_delta)
            text = text + text_delta

    return function_name, tool_call_id, args, text,return_delta


prompt = "Tell me the sum of 1, 2, 3, 4, 5, 6, 7, 8, 9, 10."
# prompt = "桌上有 2 个苹果，四个桃子和 3 本书，一共有几个水果？"
# prompt = "1+2+3...+99+100"
# prompt = "1024 乘以 1024 是多少？"   # Tools 里没有定义乘法，会怎样？
# prompt = "太阳从哪边升起？"           # 不需要算加法，会怎样？

messages = [
    {"role": "system", "content": "你是一个数学家"},
    {"role": "user", "content": prompt}
]

response = get_completion(messages)
function_name,tool_call_id, args, text,delta = getStreamResponse(response)
messages.append(delta)

while function_name != None:
    args = json.loads(args)
    result = sum123(args["numbers"])
    print("=====函数返回=====")
    print(result)

    # 把函数调用结果加入到对话历史中
    messages.append(
        {
            "tool_call_id": tool_call_id,  # 用于标识函数调用的 ID
            "role": "tool",
            "name": "sum",
            "content": str(result)  # 数值 result 必须转成字符串
        }
    )

    # 再次调用大模型
    print("=====最终回复=====")
    get_completion(messages)

    function_name, tool_call_id, args, text,delta = getStreamResponse(response)

print("=====GPT回复=====")
print_json(text)
